#!/usr/bin/python3
# -*- coding: utf-8 -*-
# 说明：自行将中英文混合文本按照中文、英文分段处理，全部使用`onnx`模式合成语音数据，然后按照原有顺序拼接起来。
import re,uuid,os
import shutil,time
from infer import TTSExecutor
from paddlespeech.resource import CommonTaskResource

executor_en = TTSExecutor()
executor_en.task_resource = CommonTaskResource(task='tts', model_format='onnx')
executor_en._init_from_path_onnx(am='fastspeech2_vctk',voc='hifigan_vctk',lang='en')

executor_zh = TTSExecutor()
executor_zh.task_resource = CommonTaskResource(task='tts', model_format='onnx')
executor_zh._init_from_path_onnx(am='fastspeech2_csmsc',voc='hifigan_csmsc',lang='zh')

def reset_executors():
	global executor_en
	global executor_zh
	executor_en = TTSExecutor()
	executor_en.task_resource = CommonTaskResource(task='tts', model_format='onnx')
	executor_en._init_from_path_onnx(am='fastspeech2_vctk',voc='hifigan_vctk',lang='en')

	executor_zh = TTSExecutor()
	executor_zh.task_resource = CommonTaskResource(task='tts', model_format='onnx')
	executor_zh._init_from_path_onnx(am='fastspeech2_csmsc',voc='hifigan_csmsc',lang='zh')
	
#spk_id = 23
spk_id = 3

def split1(s:str):
	r1 = '[a-zA-Z][a-zA-Z0-9 .,!?\'"`()（），。！“”‘’；？—]*'
	p = re.compile(r1)
	res = p.finditer(s)
	return res

def get_real_end(s:str,st:int,end:int):
	r1='[.!?\'"][ ]*\d+[ ]*$'
	p = re.compile(r1)
	v = p.search(s[st:end])
	if v==None:
		return end
	else:
		return st+v.start()+1

def preDeal(s:str):
	r1 = '[.,!\'"`()（），。！“”‘’；—]+'
	p = re.compile(r1)
	secs = p.split(s)
	return '.'.join(secs)
	

def split_zh_en(s:str):
	r2 = split1(s)
	res = []
	st = 0
	end = 0
	for v in r2:
		if v.start()>st:
			res.append( (s[st:v.start()],'zh') )
			st = v.start()
		end = get_real_end(s,st,v.end())
		res.append( (s[st:end],'en') )
		st = end
	if len(s)>end:
		res.append( (preDeal(s[end:]),'zh') )
	return res

def get_tmp():
	dname = '/tmp/'+str(uuid.uuid1())
	try:
		os.mkdir(dname)
	except:
		print('temp created')
	return dname

def concat_wav(listname:str, name:str):
	os.spawnvpe(os.P_WAIT, 'ffmpeg', 
				('ffmpeg','-f','concat','-i',listname,'-f','wav','-acodec','pcm_s16le','-y',name),
				os.environ)
	return name
	
def infer2wav(s:str, name:str):
	global spk_id
	# executor_en = TTSExecutor()
	# executor_en.task_resource = CommonTaskResource(task='tts', model_format='onnx')
	# executor_en._init_from_path_onnx(am='fastspeech2_vctk',voc='hifigan_vctk',lang='en')

	# executor_zh = TTSExecutor()
	# executor_zh.task_resource = CommonTaskResource(task='tts', model_format='onnx')
	# executor_zh._init_from_path_onnx(am='fastspeech2_csmsc',voc='hifigan_csmsc',lang='zh')

	text_seq = split_zh_en(s)
	i = 0
	tmp = get_tmp()
	listname = tmp+'/mixlist.txt'
	fp = open(listname,'w')
	for sec in text_seq:
		i+=1
		
		name1 = tmp+'/a'+str(i)+'.wav'
		name2 = 'a'+str(i)+'.wav'
		if sec[1]=='zh':
			executor_zh.infer_onnx(text=sec[0],lang='zh',am='fastspeech2_csmsc')
			executor_zh.postprocess_onnx(name1)
		elif sec[1]=='en':
			executor_en.infer_onnx(text=sec[0],lang='en',am='fastspeech2_vctk',spk_id=spk_id)
			executor_en.postprocess_onnx(name1)
		fp.write('file '+name2+'\n')
	fp.flush()
	fp.close()
	
	#os.system('ffmpeg -f concat -i %s -f wav -acodec pcm_s16le -y "%s"' % (listname,name) )
	concat_wav(listname,name)
	shutil.rmtree(tmp)
